import numpy as np

class VectorPathFollowingAlgorithm:
    def __init__(self, params):
        # 参数设置
        self.linear_speed = params['linear_speed']
        self.angular_speed_max = params['angular_speed_max']
        self.lookahead_distance = params['lookahead_distance']
        self.goal_threshold = params['goal_threshold']
        self.orientation_threshold = params['orientation_threshold']
        
        # 路径点和状态变量
        self.path = None
        self.current_waypoint_idx = 0
        self.is_following = False
    
    def set_path(self, path_msg):
        """设置要跟随的路径"""
        self.path = path_msg
        self.current_waypoint_idx = 0
        self.is_following = True
    
    def calculate_control(self, current_pose):
        """计算控制命令"""
        if not self.is_following or self.path is None:
            return 0.0, 0.0
            
        robot_x, robot_y, robot_theta = current_pose
        
        # 检查是否到达所有路径点
        if self.current_waypoint_idx >= len(self.path.poses):
            self.is_following = False
            return 0.0, 0.0
            
        # 获取当前目标点
        target_pose = self.path.poses[self.current_waypoint_idx].pose
        target_x = target_pose.position.x
        target_y = target_pose.position.y
        
        # 计算控制命令
        linear_vel, angular_vel, distance = self._calculate_vector_control(
            robot_x, robot_y, robot_theta, target_x, target_y
        )
        
        # 检查是否到达当前目标点
        if distance < self.goal_threshold:
            self.current_waypoint_idx += 1
            if self.current_waypoint_idx < len(self.path.poses):
                next_target_pose = self.path.poses[self.current_waypoint_idx].pose
                next_target_x = next_target_pose.position.x
                next_target_y = next_target_pose.position.y
                self.get_logger().info(f'到达目标点 {self.current_waypoint_idx-1}，切换到下一个目标点 ({next_target_x}, {next_target_y})')
            else:
                self.get_logger().info('已到达所有路径点')
        
        return linear_vel, angular_vel
    
    def _calculate_vector_control(self, robot_x, robot_y, robot_theta, target_x, target_y):
        """使用向量法计算控制命令"""
        # 计算目标点相对于机器人的向量
        dx = target_x - robot_x
        dy = target_y - robot_y
        
        # 计算目标距离
        distance = np.sqrt(dx*dx + dy*dy)
        
        # 计算目标角度（相对于机器人坐标系）
        target_angle = np.arctan2(dy, dx)
        angle_error = target_angle - robot_theta
        
        # 将角度误差归一化到[-pi, pi]
        while angle_error > np.pi:
            angle_error -= 2 * np.pi
        while angle_error < -np.pi:
            angle_error += 2 * np.pi
        
        # 计算线速度和角速度
        linear_vel = self.linear_speed * (1 - min(1.0, abs(angle_error) / (np.pi/2)))
        angular_vel = self.angular_speed_max * angle_error / (np.pi/2)
        
        # 限制最大线速度
        linear_vel = min(self.linear_speed, linear_vel)
        
        # 如果距离目标点很近，减速并增加角速度权重
        if distance < self.goal_threshold * 2:
            linear_vel *= distance / (self.goal_threshold * 2)
            angular_vel *= 1.5
        
        return linear_vel, angular_vel, distance